# vc_client.py
import os
import sys
import tempfile
import wave
import numpy as np
import requests
import sounddevice as sd
import soundfile as sf
from queue import Queue
from threading import Thread
from pydub import AudioSegment
from pydub.playback import play

# 配置信息
SERVER_URL = "http://localhost:8000"
AUDIO_TEMP = "output/tmp_audio.wav"
SAMPLE_RATE = 16000
CHANNELS = 1


class StreamPrinter:
    """流式文本打印器"""
    def __init__(self):
        self.print_queue = Queue()
        self.running = True
        self.thread = Thread(target=self._print_worker)
        self.thread.start()

    def _print_worker(self):
        while self.running or not self.print_queue.empty():
            try:
                chunk = self.print_queue.get(timeout=0.1)
                sys.stdout.write(chunk)
                sys.stdout.flush()
            except:
                pass

    def add_chunk(self, text):
        self.print_queue.put(text)

    def stop(self):
        self.running = False
        self.thread.join()


def record_audio(filename, duration=10):
    """重构录音函数"""
    print(f"\n开始录音（最长{duration}秒）...")

    # 使用正确的音频格式参数
    audio_data = sd.rec(
        int(duration * SAMPLE_RATE),
        samplerate=SAMPLE_RATE,
        channels=1,
        dtype='float32',
        blocking=True
    )

    # 转换为16位PCM格式
    audio_data = (audio_data * 32767).astype(np.int16)

    # 保存为WAV格式
    with wave.open(filename, 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(audio_data.tobytes())


def play_audio(data: bytes):
    """直接播放字节流音频"""
    try:
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            f.write(data)
            temp_path = f.name

        # 使用pydub播放
        audio = AudioSegment.from_wav(temp_path)
        play(audio)

        # 清理临时文件
        os.unlink(temp_path)
    except Exception as e:
        print(f"播放失败: {str(e)}")
        # 备用播放方案
        os.system(f"start {temp_path}")


def text_mode():
    printer = StreamPrinter()
    try:
        while True:
            text = input("\n请输入问题（输入q退出）: ")
            if text.lower() == 'q':
                break

            response = requests.post(
                # f"{SERVER_URL}/api/text",
                f"{SERVER_URL}/v1/stream_text",
                json={"text": text},
                stream=True,
                headers={'Accept': 'text/event-stream'}
            )

            # 处理流式响应
            for line in response.iter_lines():
                if line:
                    decoded = line.decode().strip()
                    if decoded.startswith('data:'):
                        chunk = decoded[5:].strip()
                        printer.add_chunk(chunk)
    finally:
        printer.stop()


def voice_mode():
    """重构语音交互模式"""
    try:
        while True:
            input("\n按下回车开始录音（输入q退出）...")
            record_audio(AUDIO_TEMP, duration=30)

            try:
                with open(AUDIO_TEMP, "rb") as f:
                    response = requests.post(
                        f"{SERVER_URL}/voice_to_voice",
                        files={"file": f},
                        timeout=30
                    )
                response.raise_for_status()

                # 获取完整音频
                audio_data = response.content

                print("\n正在播放回答...")
                play_audio(audio_data)

            except Exception as e:
                print(f"\n连接错误: {str(e)}")

    except KeyboardInterrupt:
        pass


def main():
    """主界面"""
    os.makedirs(os.path.dirname(AUDIO_TEMP), exist_ok=True)

    print("""\n语音交互测试客户端
    1. 文本模式
    2. 语音模式
    3. 退出""")

    while True:
        choice = input("\n请选择模式: ")
        if choice == '1':
            text_mode()
        elif choice == '2':
            voice_mode()
        elif choice == '3':
            break
        else:
            print("无效输入")


if __name__ == "__main__":
    main()
